import mysql.connector
from sqlalchemy import create_engine
import pandas as pd
import logging
from utils.redis_client import RedisClient

class DatabaseConnection:
    """
    数据库连接公共类
    提供统一的数据库连接和操作接口，集成Redis缓存
    """

    def __init__(self, db_config, redis_client=None):
        """
        初始化数据库连接配置
        :param db_config: 数据库配置字典
        :param redis_client: Redis客户端实例（可选）
        """
        self.db_config = db_config
        self.redis_client = redis_client
        self.connection = None

    def get_mysql_connector_connection(self):
        """
        获取mysql.connector连接对象
        :return: mysql.connector连接对象
        """
        try:
            connection = mysql.connector.connect(**self.db_config)
            return connection
        except Exception as e:
            logging.error(f"创建mysql.connector连接失败: {str(e)}")
            raise

    def get_sqlalchemy_engine(self):
        """
        获取SQLAlchemy引擎对象
        :return: SQLAlchemy引擎对象
        """
        try:
            engine_url = f"mysql+mysqlconnector://{self.db_config['user']}:{self.db_config['password']}@{self.db_config['host']}/{self.db_config['database']}"
            engine = create_engine(engine_url)
            return engine
        except Exception as e:
            logging.error(f"创建SQLAlchemy引擎失败: {str(e)}")
            raise

    def execute_query(self, query, cache_key=None, cache_expire=3600):
        """
        执行查询语句并返回结果
        :param query: SQL查询语句
        :param cache_key: Redis缓存键（可选）
        :param cache_expire: 缓存过期时间（秒）
        :return: 查询结果DataFrame
        """
        try:
            # 如果启用了Redis缓存且提供了缓存键
            if self.redis_client and cache_key:
                # 尝试从缓存获取数据
                cached_data = self.redis_client.get(cache_key)
                if cached_data is not None:
                    logging.info(f"从Redis缓存获取数据: {cache_key}")
                    return cached_data

            # 从数据库查询数据
            engine = self.get_sqlalchemy_engine()
            data = pd.read_sql(query, engine)

            # 如果启用了Redis缓存且提供了缓存键，将结果存入缓存
            if self.redis_client and cache_key:
                self.redis_client.set(cache_key, data, ex=cache_expire)
                logging.info(f"将查询结果存入Redis缓存: {cache_key}")

            return data
        except Exception as e:
            logging.error(f"执行查询失败: {str(e)}")
            raise

    def execute_update(self, query, params=None):
        """
        执行更新语句（INSERT, UPDATE, DELETE）
        :param query: SQL更新语句
        :param params: 参数列表
        :return: 影响的行数
        """
        connection = None
        cursor = None
        try:
            connection = self.get_mysql_connector_connection()
            cursor = connection.cursor()
            if params:
                cursor.execute(query, params)
            else:
                cursor.execute(query)
            connection.commit()
            return cursor.rowcount
        except Exception as e:
            if connection:
                connection.rollback()
            logging.error(f"执行更新失败: {str(e)}")
            raise
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()

    def close_all_connections(self):
        """
        关闭所有连接（占位方法，实际连接使用后即关闭）
        """
        pass
